import os
import random
import re
import math
import struct
import logging
import datetime
import pandas as pd
from collections import namedtuple
from pandas import Timedelta
from pytdx.config.hosts import hq_hosts
from pytdx.hq import TdxHq_API
from pytdx.params import TDXParams
from pytdx.pool.hqpool import TdxHqPool_API
from pytdx.pool.ippool import AvailableIPPool
from pytdx.reader import TdxLCMinBarReader, TdxFileNotFoundException, BlockReader
from pytdx.reader.block_reader import BlockReader_TYPE_GROUP

from pyLibs import GadflyUnits as units
from apscheduler.schedulers.blocking import BlockingScheduler
from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.events import EVENT_JOB_ERROR, EVENT_JOB_EXECUTED

logging.basicConfig(format="%(asctime)s %(filename)s(line:%(lineno)d) [%(levelname)s] : %(message)s", datefmt="%Y-%M-%d %H:%M:%S", level=logging.DEBUG)


pd.set_option('display.max_columns', None)  # 显示所有列
pd.set_option('display.max_rows', None)
pd.set_option('expand_frame_repr', False)  # 显示宽度
pd.set_option('display.unicode.ambiguous_as_wide', True)
pd.set_option('display.unicode.east_asian_width', True)
#pd.set_option('display.width', 180) # 设置打印宽度(**重要**)
pd.set_option('display.unicode.east_asian_width', True) #设置输出右对齐


'''通过证券代码获取市场信息'''


def get_stock_market(code):
    patterns = (
        (0, '深圳指数', '深圳证券交易所', '^399\d{3}$', 'sz'),
        (1, '上证A股', '上海证券交易所', '^60[0135]\d{3}$', 'sh'),
        (0, '深证A股', '深圳证券交易所', '^00\d{4}$', 'sz'),
        (1, '上证B股', '上海证券交易所', '^900\d{3}$', 'sh'),
        (0, '深证B股', '深圳证券交易所', '^200\d{3}$', 'sz'),
        (2, '北证A股', '北京证券交易所', '^(43|83|87|88)\d{4}$', 'bj'),
        (1, '中小板', '深圳证券交易所', '^002\d{3}$', 'sz'),
        (1, '科创板', '上海证券交易所', '^68\d{4}$', 'sh'),
        (0, '创业板', '深圳证券交易所', '^30\d{4}$', 'sz'),
        (3, '港股', '香港证券交易所', '\d{4,5}.HK$', 'hk'),
        (1, '上证国债现货', '上海证券交易所', '^001\d{3}$', 'sh'),
        (1, '上证企业债券', '上海证券交易所', '^0[12]0\d{3}$', 'sh'),
        (1, '上证可转换债券', '上海证券交易所', '^1[20][90]\d{3}$', 'sh'),
        (1, '上证国债回购', '上海证券交易所', '^204\d{3}$', 'sh'),
        (1, '上证国债期货', '上海证券交易所', '^310\d{3}$', 'sh'),
        (0, '深证国债现货', '深圳证券交易所', '^10\d{4}$', 'sz'),
        (0, '深证企业债券', '深圳证券交易所', '^11\d{4}$', 'sz'),
        (0, '深证可转换债券', '深圳证券交易所', '^12\d{4}$', 'sz'),
        (0, '深证国债回购', '深圳证券交易所', '^131\d{3}$', 'sz'),
        (1, '上证封闭式基金', '上海证券交易所', '^508\d{3}$', 'sh'),
        (1, '上证ETF基金', '上海证券交易所', '^(51|56|58)\d{4}$', 'sh'),
        (1, '上证LOF基金', '上海证券交易所', '^501\d{3}$', 'sh'),
        (1, '上证分级基金', '上海证券交易所', '^(502|506)\d{3}$', 'sh'),
        (0, '深证封闭式基金', '深圳证券交易所', '^184\d{3}$', 'sz'),
        (0, '深证ETF基金', '深圳证券交易所', '^159\d{3}$', 'sz'),
        (0, '深证LOF基金', '深圳证券交易所', '^16\d{4}$', 'sz'),
        (0, '深证分级基金', '深圳证券交易所', '^15\d{4}$', 'sz'),
    )
    for m in patterns:
        if re.match(m[3], code):
            return {
                'stock_code': code,
                'market': m[0],
                'stock_securities': m[1],
                'stock_market': m[2],
                'ths_code': code + '.' + m[4],
                'tdx_code': m[4] + code
            }


'''获取通信达终端的板块信息'''


def get_stock_block(blockname, type='default', tdx_path='D:/通达信金融终端'):
    block_path = os.path.join(os.path.dirname(tdx_path), 'T0002', 'hq_cache')
    if type == 'zs':
        block_file = os.path.join(block_path, 'block_zs.dat')
    elif type == 'fg':
        block_file = os.path.join(block_path, 'block_fg.dat')
    elif type == 'gn':
        block_file = os.path.join(block_path, 'block_gn.dat')
    else:
        block_file = os.path.join(block_path, 'block.dat')
    bf = BlockReader().get_df(block_file, BlockReader_TYPE_GROUP)
    if bf is None:
        return bf
    else:
        if blockname == None:
            dateList = bf['code_list'].values[0].split(',')
        else:
            if isinstance(blockname, str):
                dateList = bf.query('blockname == ' + blockname)['code_list'].values[0].split(',')
            elif isinstance(blockname, (tuple, list)):
                if isinstance(blockname, tuple):
                    blockname = list(blockname)
                dateList = ','.join(bf.query('blockname in ' + str(blockname))['code_list'].values).split(',')
            else:
                dateList = bf['code_list'].values[0].split(',')
        return pd.Series(dateList)


class Scheduler:
    stockData = None

    def __init__(self):
        self.tdxStore = TdxStore()
        self.scheduler = BlockingScheduler()
        #self.scheduler = BackgroundScheduler()

    def live_connect(self, code):
        self.scheduler.add_listener(self.listener, EVENT_JOB_EXECUTED | EVENT_JOB_ERROR)
        # 添加任务并设置触发方式为3s一次
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='9',
                               minute='35, 40, 50, 55', second='2', args=[code, '5m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='9',
                               minute='45', second='2', args=[code, '15m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='11',
                               minute='05, 10, 20, 25', second='2', args=[code, '5m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='10',
                               minute='00, 15, 45', second='2',  args=[code, '15m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='11',
                               minute='00, 15', second='2', args=[code, '15m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='10-11',
                               minute='30', second='2', args=[code, '1h'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='10, 13-15',
                               minute='05, 15, 20, 25, 35, 40, 50, 55', second='2', args=[code, '5m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='13-15',
                               minute='15, 30, 45', second='2', args=[code, '15m'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='17',
                               minute='01', second='2', args=[code, '1h'])
        self.scheduler.add_job(self.job_executes, 'cron', day_of_week='mon-fri', hour='09-23',
                               minute='00-59', second='2, 32', args=[code, '5m'])
        self.scheduler.start()

    @classmethod
    def getStockData(cls):
        return cls.stockData

    def listener(self, event):
        if event.exception:
            print('The job did not run')
        else:
            self.tdxStore.onStockBar(self.stockData)

    def job_executes(self, code, rule='5m'):
        logging.info("开始执行任务，任务参数：%s" % rule)
        if rule in ['1m', '1M', '1t', '1T', 8]:
            category = 8
        elif rule in ['5m', '5M', '5t', '5T', 0]:
            category = 0
        elif rule in ['15m', '15M', '15t', '15T', 1]:
            category = 1
        elif rule in ['30m', '30M', '30t', '30T', 2]:
            category = 2
        elif rule in ['60m', '60M', '60t', '60T', '1h', '1H', 3]:
            category = 3
        elif rule in ['1d', '1D', 'Day', 4]:
            category = 4
        elif rule in ['1w', '1W', 'Week', 5]:
            category = 5
        elif rule in ['month', 'Month', 'MM', 6]:
            category = 6
        elif rule in ['quarter', 'Quarter', '1Q', '1q', 10]:
            category = 10
        elif rule in ['year', 'Year', '1Y', '1q', 11]:
            category = 11
        else:
            category = 9
        self.stockData = self.tdxStore.getStockBarsFromApi(8, code, 0, 2, True)
        logging.info("执行任务完成，任务参数：%s" % rule)


class TdxStore:
    def __init__(self, param=None):
        if param == None:
            self.TdxPath = param
        else:
            config = units.loadYaml('config.yaml')
            platform = units.checkPlatform()
            if platform == 'Windows':
                self.TdxPath = config.stock_source.tdx.win_path
            elif platform in ['MacOS', 'Linux', 'Windows Subsystem for Linux (WSL)']:
                self.TdxPath = config.stock_source.tdx.mac_path
            else:
                pass

    def onStockBar(self, data):
        print(data)
        return data

    def get_stock_data(self, code, cycle=None):
        info = get_stock_market(code)
        local_data = self.getStockBarsFromLocal(info['ths_code'], cycle)
        if cycle in ['1d', '1D', 'd', 'D']:
            network_data = self.getStockBarsFromApi(4, info['market'], code, 0, 30, True)
        elif cycle in ['5m', '5M', '5t', '5T']:
            network_data = self.getStockBarsFromApi(0, info['market'], code, 0, 600, True)
        elif cycle in ['15m', '15M', '15t', '15T']:
            local_data = self.lc5Resample(local_data, rule='15T')
            network_data = self.getStockBarsFromApi(1, info['market'], code, 0, 400, True)
        elif cycle in ['30m', '30M', '30t', '30T']:
            local_data = self.lc5Resample(local_data, rule='30T')
            network_data = self.getStockBarsFromApi(2, info['market'], code, 0, 200, True)
        elif cycle in ['60m', '60M', '60t', '60T', '1h', '1H', 'h', 'H']:
            local_data = self.lc5Resample(local_data, rule='60T')
            network_data = self.getStockBarsFromApi(3, info['market'], code, 0, 100, True)
        else:
            network_data = pd.DataFrame()
        stock_data = pd.concat([local_data, network_data], ignore_index=True)
        stock_data = stock_data.drop_duplicates(subset=['datetime'], keep='first', ignore_index=True)
        stock_data.index = pd.to_datetime(stock_data.datetime)
        Stock = namedtuple('Stock', ['code', 'data', 'length', 'columns', 'start', 'end'])
        return Stock(info['tdx_code'], stock_data, stock_data.shape[0], stock_data.columns,
                     stock_data['datetime'][0], stock_data["datetime"].tail(1).values[0])

    def resolve(self, *args, **kwargs):
        return os.path.join(self.TdxPath, *args, **kwargs)

    '''
    获取股票k线数据
    * category-> K线种类
        0 5分钟K线 
        1 15分钟K线 
        2 30分钟K线 
        3 1小时K线 
        4 日K线
        5 周K线
        6 月K线
        7 1分钟
        8 1分钟K线 9 日K线
        10 季K线
        11 年K线
    * market -> 市场代码 0:深圳，1:上海
    * stockcode -> 证券代码;
    * start -> 指定的范围开始位置;
    * count -> 用户要请求的 K 线数目，最大值为 800
    '''

    def getStockBarsFromApi(self, category, code, start, count, isDp=False):
        info = get_stock_market(code)
        api = TdxHq_API()
        if api.connect("119.147.212.81", 7709):
            if (isDp):
                data = api.to_df(api.get_security_bars(category, info['market'], code, start, count))
                df_data = pd.DataFrame()
                df_data = df_data.assign(datetime=data['datetime'])
                df_data = df_data.assign(
                    date=data.apply(lambda x: datetime.date(x['year'], x['month'], x['day']), axis=1))
                df_data = df_data.assign(time=data.apply(lambda x: datetime.time(x['hour'], x['minute']), axis=1))
                df_data = df_data.assign(open=data['open'])
                df_data = df_data.assign(high=data['high'])
                df_data = df_data.assign(low=data['low'])
                df_data = df_data.assign(close=data['close'])
                df_data = df_data.assign(volume=data['vol'])
                df_data = df_data.assign(ratio=data['close'].diff())
                df_data = df_data.assign(amount=data['amount'])
                df_data = df_data.assign(openinterest=0)
                df_data = df_data.dropna()  # 去除空值
                df_data.index = pd.to_datetime(df_data.datetime)
                return df_data
            else:
                return api.get_security_bars(category, info['market'], code, start, count)

    # 获取股票指数数据
    '''
    * category-> K线种类
        0 5分钟K线 
        1 15分钟K线 
        2 30分钟K线 
        3 1小时K线 
        4 日K线
        5 周K线
        6 月K线
        7 1分钟
        8 1分钟K线 
        9 日K线
        10 季K线
        11 年K线
    * market -> 市场代码 0:深圳，1:上海
    * stockcode -> 证券代码;
    * start -> 指定的范围开始位置;
    * count -> 用户要请求的 K 线数目，最大值为 800
    '''

    def getIndexBarsFromApi(self, category, code, start, count, isDp=False):
        info = get_stock_market(code)
        ips = [(v[1], v[2]) for v in hq_hosts]
        random.shuffle(ips)  # 获取5个随机ip作为ip池
        ips5 = ips[:5]
        ippool = AvailableIPPool(TdxHq_API, ips5)  ## IP 池对象
        primary_ip, hot_backup_ip = ippool.sync_get_top_n(2)  ## 选出M, H
        ## 生成hqpool对象，第一个参数为TdxHq_API后者 TdxExHq_API里的一个，第二个参数为ip池对象。
        api = TdxHqPool_API(TdxHq_API, ippool)
        ## connect 函数的参数为M, H 两组 (ip, port) 元组
        with api.connect(primary_ip, hot_backup_ip):
            if (isDp):
                return api.to_df(api.get_index_bars(category, info['market'], code, start, count))
            else:
                return api.get_index_bars(category, info['market'], code, start, count)

    # 根据通达信5分钟周期数据，生成其他周期数据
    def lc5Resample(self, data_lc5, rule):
        if rule == '60T':
            data_30 = data_lc5.resample('30T', closed='right', label='right').agg({
                'datetime': 'last', 'date': 'last', 'time': 'last', 'open': 'first', 'high': 'max',
                'low': 'min', 'close': 'last', 'volume': 'sum', 'amount': 'sum'}).dropna()
            if data_30.shape[0] % 2 != 0:
                raise Exception("30分钟数据合并时出现错误。Data_30T：%d" % len(data_30))
            # 使用 GroupBy 构造60分钟 K线
            idx = list(range(0, len(data_30) // 2)) * 2
            idx.sort()
            data_30['idx'] = idx
            data_30 = data_30.reset_index()
            df_data = data_30.groupby(by='idx').agg({
                'datetime': 'last', 'date': 'last', 'time': 'last', 'open': 'first', 'high': 'max',
                'low': 'min', 'close': 'last', 'volume': 'sum', 'amount': 'sum'}).dropna()
        else:
            df_data = data_lc5.resample(rule, closed='right', label='right').agg({
                'datetime': 'last', 'date': 'last', 'time': 'last', 'open': 'first', 'high': 'max',
                'low': 'min', 'close': 'last', 'volume': 'sum', 'amount': 'sum'}).dropna()
        df_data = df_data.assign(ratio=df_data['close'].diff())
        df_data = df_data.assign(openinterest=0)
        df_data.index = pd.to_datetime(df_data.datetime)
        return df_data

    # 获取本地K线数据
    def getStockBarsFromLocal(self, ts_code, cycle=None):
        if (cycle in ('1h', '60m', '30m', '15m', '5m')):
            filename = ts_code.split(".")[1] + ts_code.split(".")[0] + ".lc5"
            filepath = self.resolve("vipdoc", ts_code.split(".")[1], "fzline", filename)
            pack = 'HHffffllf'
        elif (cycle == '1m'):
            filename = ts_code.split(".")[1] + ts_code.split(".")[0] + ".lc1"
            filepath = self.resolve("vipdoc", ts_code.split(".")[1], "minline", filename)
            pack = 'HHffffllf'
        else:
            filename = ts_code.split(".")[1] + ts_code.split(".")[0] + ".day"
            filepath = self.resolve("vipdoc", ts_code.split(".")[1], "lday", filename)
            pack = 'lllllfll'
        kline = self._read_kline(filepath, pack)
        # kline["ts_code"] = ts_code
        kline.index = pd.to_datetime(kline["datetime"])
        kline.index.name = "index"
        # kline = kline.rename(columns={"vol": "volume"})
        usecols = ("datetime date time open high low close ratio amount volume openinterest".split())
        return kline[usecols]

    # 根据二进制前两段拿到日期分时
    def _parseDateTime(self, h1, h2) -> str:  # H1->0,1字节; H2->2,3字节;
        return self._parseDate(h1) + " " + self._parseTime(h2)

    def _formatDate(self, year, m, d):
        if (m < 10):
            month = '0' + str(m)
        else:
            month = str(m)
        if (d < 10):
            day = '0' + str(d)
        else:
            day = str(d)
        return str(year) + "-" + month + "-" + day

    def _parseDate(self, h):
        year = math.floor(h / 2048) + 2004  # 解析出年
        month = math.floor(h % 2048 / 100)  # 月
        day = h % 2048 % 100  # 日
        return self._formatDate(year, month, day)

    def _parseTime(self, h):
        hour = math.floor(h / 60)  # 小时
        minute = h % 60  # 分钟
        if hour < 10:  # 如果小时小于两位, 补0
            hour = "0" + str(hour)
        if minute < 10:  # 如果分钟小于两位, 补0
            minute = "0" + str(minute)
        return str(hour) + ":" + str(minute)

    # 读取K线源文件
    def _read_kline(self, filepath, pack='lllllfll'):
        with open(filepath, "rb") as f:
            usecols = "datetime date time open high low close ratio amount volume openinterest".split()
            buffers = []
            i = 0
            while True:
                buffer = f.read(32)
                if not buffer:
                    break
                buffer = struct.unpack(pack, buffer)
                if i == 0:
                    preClose = buffer[4] / 100
                ratio = round((buffer[4] / 100 - preClose) / preClose * 100, 3)
                preClose = buffer[4] / 100
                i = i + 1
                if (pack == 'lllllfll'):
                    year = int(buffer[0] / 10000)
                    month = int((buffer[0] % 10000) / 100)
                    day = (buffer[0] % 10000) % 100
                    date = self._formatDate(year, month, day)
                    buffers.append((date + ' 15:00', date, '15:00', buffer[1] / 100, buffer[2] / 100, buffer[3] / 100,
                                    buffer[4] / 100, ratio, buffer[5], buffer[6], buffer[7]))
                else:
                    buffers.append((self._parseDateTime(buffer[0], buffer[1]), self._parseDate(buffer[0]),
                                    self._parseTime(buffer[1]),
                                    buffer[2], buffer[3], buffer[4], buffer[5], ratio, buffer[6], buffer[7], buffer[8]))
            kline = pd.DataFrame(buffers, columns=usecols)
        kline["datetime"] = kline["datetime"].astype(str)
        price_columns = ["open", "high", "low", "close"]
        kline[price_columns] = kline[price_columns].apply(lambda x: x)
        return kline
